#include "haptree.h"
#include<math.h>

ReadEvidenceList *allocNewEvidence(ReadEvidence *rEvidence){
  ReadEvidenceList * newEvidence = malloc(sizeof(ReadEvidenceList));
  newEvidence->rEvidence = rEvidence;
  return newEvidence;
}

void addNodeEvidence(HapNode * node, ReadEvidenceList * nodeEvidence){
  ReadEvidenceList *temp;
  
  if(nodeEvidence->offset == nodeEvidence->rEvidence->pRead->length){
    nodeEvidence->rEvidence->pRead->refCount-=1;
    nodeEvidence->offset = 0;
    if(nodeEvidence->rEvidence->pRead->refCount==0){
      freeRead(nodeEvidence->rEvidence->pRead);
      nodeEvidence->rEvidence->pRead = NULL;
    }
  }

  if(nodeEvidence->offset != 0 || node->evidenceList == NULL){
    nodeEvidence->next = node->evidenceList;
    node->evidenceList = nodeEvidence;
  }
  else{
    for(temp=node->evidenceList; temp->next!=NULL; temp=temp->next){
      if(temp->next->offset == 0){
	nodeEvidence->next = temp->next;
	temp->next = nodeEvidence;
	return;
      }
    }
    temp->next = nodeEvidence;
    nodeEvidence->next = NULL;
  }
}

void extendIncompleteEvidence(HaplotypeCollection * hapCol){
  HapNode * parent;
  ReadEvidenceList *eviL;
  if(hapCol->currentLoci==1)
    return;

  for(parent = hapCol->loci[hapCol->currentLoci-2]->allNodes;
      parent != NULL; parent=parent->next){
    while(parent->evidenceList!=NULL && parent->evidenceList->offset != 0){
      eviL = parent->evidenceList;
      parent->evidenceList=eviL->next;
      eviL->rEvidence->pRead->refCount +=matchPath(parent,eviL)-1;
    }
  }
}

/* HapNode */
HapNode *newHapNode(unsigned int allele){
  HapNode *node=malloc(sizeof(HapNode));
  node->allele = allele;
  node->next = NULL;
  node->next_sibling = NULL;
  node->prev_sibling = NULL;
  node->children = NULL;
  node->parent = NULL;
  node->evidenceList = NULL;
  node->totalProb = 0.0;
  node->nChildren = 0;
  return node;
}

void freeHapNode(HapNode * node){
  ReadEvidenceList *tempEvi;
  if(node!=NULL)
    while(node->evidenceList !=NULL){
      tempEvi=node->evidenceList;
      node->evidenceList = tempEvi->next;
      free(tempEvi);
    }
    free(node);
}


/* Locus */
Locus *allocateLocus(unsigned int length, unsigned int nAlleles){
  Locus * locus = malloc(sizeof(Locus));
  unsigned int allele, aHap;
  HapNode *prev, *current;

  locus->nAlleles = nAlleles;
  locus->alleleIndex = malloc(sizeof(HapNode *)*nAlleles);
  locus->totalNodes = length*nAlleles;
  prev = NULL;
  for(allele=0; allele<nAlleles; ++allele){
    for(aHap=0; aHap<length; ++aHap){
      current = newHapNode(allele);
      current->totalProb = 1.0/nAlleles;
      if(prev!=NULL){
	prev->next=current;
      }
      else{
	locus->allNodes = current;
      }
      if(aHap==0){
	locus->alleleIndex[allele] = current;
      }
      prev=current;
    }
  }
  return locus;
}


void freeLocus(Locus *locus){
  if(locus!=NULL){
    if(locus->alleleIndex!=NULL)
      free(locus->alleleIndex);
    free(locus);
  }
}

Locus * addLocus(Locus * lastLocus, unsigned int nAlleles){
  Locus * locus;
  locus = allocateLocus(lastLocus->totalNodes, nAlleles);
  linkNextLocus(lastLocus, locus);

  return locus;
}

void linkNextLocus(Locus *lastLocus, Locus *newLocus){
  HapNode *parentNode;
  HapNode *prev;
  HapNode **parentAllele;
  unsigned int allele;
  double prob;
  parentAllele = malloc(sizeof(HapNode *)*newLocus->nAlleles);

  for(allele = 0; allele<newLocus->nAlleles; ++allele){
    parentAllele[allele] = newLocus->alleleIndex[allele];
  }

  for(parentNode = lastLocus->allNodes; 
      parentNode!=NULL; parentNode=parentNode->next){
    parentNode->nChildren = newLocus->nAlleles;
    parentNode->children = parentAllele[0];
    prev=parentAllele[0];
    prev->parent = parentNode;
    parentAllele[0] = parentAllele[0]->next;
    prev->totalProb = prob = parentNode->totalProb/newLocus->nAlleles;

    for(allele = 1; allele<newLocus->nAlleles; ++allele){
      prev->next_sibling = parentAllele[allele];
      prev->next_sibling->prev_sibling = prev;
      prev=prev->next_sibling;
      prev->parent = parentNode;
      prev->totalProb = prob;
      parentAllele[allele] = parentAllele[allele]->next;
    }
  }
  free(parentAllele);
}

void extendHaplotypes(HaplotypeCollection * hapCol, unsigned int nAlleles){
  Locus *newLocus;
  if(hapCol->currentLoci == 0){
    newLocus = allocateLocus(1,nAlleles);
  }
  else{
    newLocus = addLocus(hapCol->loci[hapCol->currentLoci-1], nAlleles);
  }
  
  hapCol->loci[hapCol->currentLoci] = newLocus;
  hapCol->currentLoci++;
  extendIncompleteEvidence(hapCol);
}

void rebuildAlleleIndex(HaplotypeCollection *hapCol){
  Locus *locus;
  unsigned int i, lastAllele;
  HapNode *node;

  for(i = 0; i<hapCol->currentLoci; ++i){
    locus = hapCol->loci[i];
    lastAllele = locus->nAlleles+1;
    for(node = locus->allNodes; node!=NULL; node=node->next){
      if(node->allele!=lastAllele){
	locus->alleleIndex[node->allele] = node;
      }
    }
  }

}

void pruneHaplotypes(HaplotypeCollection *hapCol, double significance){
  HapNode *haplotype, *prev, *temp;
  Locus *lastLocus;

  lastLocus = hapCol->loci[hapCol->currentLoci-1];
  prev=NULL;
  haplotype = lastLocus->allNodes;
  while(haplotype!=NULL){
    temp=haplotype->next;
    if(haplotype->totalProb < significance){
      deleteHaplotype(hapCol, haplotype, prev, hapCol->currentLoci-1);
      haplotype=temp;
    }
    else{
      prev=haplotype;
      haplotype=haplotype->next;
    }
  }
}

void setHaplotypeFreqs(HaplotypeCollection *hapCol, double *freqs){
  HapNode *haplotype;
  Locus *lastLocus;  
  double ff;
  lastLocus = hapCol->loci[hapCol->currentLoci-1];
  haplotype = lastLocus->allNodes;  
  ff= 1.0/lastLocus->totalNodes;

  if(freqs==NULL){
    freqs = &ff;
  }
  while(haplotype!=NULL){
    haplotype->totalProb = *freqs;
    haplotype = haplotype->next;
    if(freqs!=&ff){
      freqs++;
    }
  }
}

void computeInternalProbs(HaplotypeCollection *hapCol){
  Locus *locus;
  unsigned int locusI;
  HapNode *node, *child;
  double pSum;

  if(hapCol->currentLoci<2){
    return;
  }

  for(locusI = hapCol->currentLoci-1; locusI>0; locusI--){
    locus = hapCol->loci[locusI-1];
    for(node = locus->allNodes; node!=NULL; node=node->next){
      pSum=0;
      for(child = node->children; child!=NULL; child=child->next_sibling){
	pSum += child->totalProb;
      }
      node->totalProb=pSum;
    }
  }

}

void computeEvidenceWeights(HaplotypeCollection *hapCol){
  /* Evidence weights are b_{r}'s in the paper */
  ReadEvidence *evi;
  ReadEvidenceList *eviL;
  unsigned int locusI;
  double tProb;
  Locus *locus;
  HapNode *node;
  /* Set all read weights to zero*/
  for(evi = hapCol->evidenceCollection.next; evi!=NULL; evi=evi->next){
    evi->pWeight = 0;
  }

  for(locusI=0; locusI < hapCol->currentLoci; locusI++){
    locus = hapCol->loci[locusI];
    for(node=locus->allNodes; node!=NULL; node=node->next){
      tProb = node->totalProb;
      for(eviL = node->evidenceList; eviL!=NULL; eviL=eviL->next){
	eviL->rEvidence->pWeight+=tProb;
      }
    }
  }

  for(evi = hapCol->evidenceCollection.next; evi!=NULL; evi=evi->next){
    evi->pWeight = evi->multiplicity/evi->pWeight;
  }

}

void computeUpdatedProbs(HaplotypeCollection *hapCol){
  unsigned int locusI;
  Locus *locus;
  HapNode *node;
  ReadEvidenceList * eviL;
  double pSum;

  for(locusI=0; locusI < hapCol->currentLoci-1; locusI++){
    locus = hapCol->loci[locusI];
    for(node=locus->allNodes; node!=NULL; node=node->next){
      if(node->parent!=NULL){
	node->totalProb=node->parent->totalProb;
      }else{
	node->totalProb=0.0;
      }
      for(eviL=node->evidenceList; eviL!=NULL; eviL=eviL->next){
	node->totalProb+=eviL->rEvidence->pWeight;
      }
    }
  }

  locus = hapCol->loci[hapCol->currentLoci-1];
  for(node=locus->allNodes; node!=NULL; node=node->next){
    pSum=node->parent->totalProb;
    for(eviL = node->evidenceList; eviL!=NULL; eviL=eviL->next){
      pSum += eviL->rEvidence->pWeight;
    }
    node->totalProb *= pSum;
  }
}

double computeCollectionFreqs(HaplotypeCollection *hapCol, Parameters *pars, double *freqs){
  unsigned int i;
  HapNode *node;
  double norm1, norm2, d, nmo2, npo2, nmo4, div;
  int clearFreqs=0;

  if(hapCol->currentLoci < 2){
    return 0;
  }

  if(freqs==NULL){
    clearFreqs=1;
    freqs = malloc(sizeof(double)*hapCol->loci[hapCol->currentLoci-1]->totalNodes);
  }
  /* copy current probabilities */
  i=0;
  for(node=hapCol->loci[hapCol->currentLoci-1]->allNodes; node!=NULL; node=node->next){
    freqs[i] = node->totalProb;
    i++;
  }

  computeInternalProbs(hapCol);
  //printf("COMPUTING INTERNAL PROBS\n\n");
  //printCollection(hapCol);
  computeEvidenceWeights(hapCol);
  //printf("\n\nCOMPUTING EVIDENCE WEIGHTS\n\n");
  //printCollection(hapCol);
  computeUpdatedProbs(hapCol);
  //printf("\n\nCOMPUTING UNNORMALIZED PROBS\n\n");
  //printCollection(hapCol);

  /* Compute normalization */
  i=0;
  d=0;
  for(node=hapCol->loci[hapCol->currentLoci-1]->allNodes; node!=NULL; node=node->next){
    if(freqs[i]>=pars->delta){
      d+=node->totalProb;
    }
    i++;
  }
  //printf("d=%0.4f\n", d);
  if(d!=0){
    npo2 = (hapCol->N + pars->omega)/2.0;
    nmo2 = npo2 - pars->omega; //(hapCol->N - pars->omega)/2.0;
    nmo4 = sqrt(pow(nmo2,2.0)+d*pars->omega);

    norm1 = nmo2 + nmo4;
    norm2 = npo2 + nmo4;
  }
  else{
    norm1 = norm2 = hapCol->N;
  }
  //printf("npo2 = %0.4f, nmo2 = %0.4f, nmo4 = %0.4f \n", npo2, nmo2, nmo4);

  //printf("norm1 = %0.4f, norm2 = %0.4f\n", norm1, norm2);
  //printf("N = %4.0f, omega = %0.4f\n", hapCol->N, pars->omega);
  /* finally normalize the frequencies */
  i=0;
  div=0.0;
  for(node=hapCol->loci[hapCol->currentLoci-1]->allNodes; node!=NULL; node=node->next){
    if(freqs[i]>=pars->delta){
      node->totalProb/=norm1;
    }
    else{
      node->totalProb/=norm2;
    }
    div += fabs(freqs[i] - node->totalProb);
    i++;
  }

  if(clearFreqs){
    free(freqs);
  }

  return div;
}

void deleteHaplotype(HaplotypeCollection *hapCol, HapNode *node, 
		     HapNode *prevNode, unsigned int locus){
  HapNode * parent = node->parent;
  HapNode *temp;
  unsigned int i;

  hapCol->loci[locus]->totalNodes--;
  if(prevNode == NULL){
    hapCol->loci[locus]->allNodes = node->next;
  }
  else{
    prevNode->next = node->next;
  }
  
  if(node == hapCol->loci[locus]->alleleIndex[node->allele]){
    if(node->next!=NULL && node->next->allele == node->allele)
      hapCol->loci[locus]->alleleIndex[node->allele] = node->next;
    else
      hapCol->loci[locus]->alleleIndex[node->allele] = NULL;
  }

  if(parent != NULL){
    
    if(node == parent->children){
      parent->children = node->next_sibling;
    }
    
    if(node->prev_sibling !=NULL){
      node->prev_sibling->next_sibling = node->next_sibling;
    }
    
    if(node->next_sibling != NULL){
      node->next_sibling->prev_sibling = node->prev_sibling;
    }
    
    parent->nChildren--;
    if(parent->nChildren == 0){
      if(parent == hapCol->loci[locus-1]->allNodes){
	deleteHaplotype(hapCol, parent, NULL, locus-1);
      }
      else{
	temp = hapCol->loci[locus-1]->alleleIndex[parent->allele];
	if(temp == parent){
	  temp=NULL;
	  for(i=1; temp==NULL; ++i){
	    temp = hapCol->loci[locus-1]->alleleIndex[parent->allele-i];
	  }
	}
	while(temp->next != parent){
	  temp=temp->next;
	}
	deleteHaplotype(hapCol, parent, temp, locus-1);
      }
    }
  }

  freeHapNode(node);  
}

void printCollection(HaplotypeCollection *hapCol){
  HapNode *node;
  if(hapCol->currentLoci == 0){
    printf("Empty haplotype collection over %zu loci\n", hapCol->totalLoci);
    return;
  }
  printf("Printing haplotype collection\n%zu loci\n%zu loci expanded\n%zu haplotypes in total\n%zu reads for evidence", hapCol->totalLoci, hapCol->currentLoci, 
	 hapCol->loci[hapCol->currentLoci-1]->totalNodes, hapCol->N);
  printf("Haplotypes:\n");
  
  for(node=hapCol->loci[0]->allNodes; node!=NULL; node=node->next){
    printHaplotype(node, 0, 0);
  }
}

void printHaplotype(HapNode *node, unsigned int locus, unsigned int skip){
  HapNode *child;
  unsigned int i;
  unsigned int nReads = 0;
  ReadEvidenceList *temp;

  if(node==NULL){
    printf(".\t\t");
    return;
  }

  if(skip){
    for(i=0;i<locus; ++i){
      printf("^\t\t");
    }
  }
  
  for(temp = node->evidenceList; temp!=NULL; temp=temp->next){
    nReads = nReads + temp->rEvidence->multiplicity;
  }

  printf("%zu{%zu}:%1.4f\t", node->allele, nReads,node->totalProb);

  if(node->nChildren ==0){
    printf("\n");
  }
  else{
    child=node->children;
    printHaplotype(child, locus+1, 0);
    for(child=child->next_sibling; child!=NULL; child=child->next_sibling){
      printHaplotype(child, locus+1, 1);
    }
  }
}


double getListHaplotype(HaplotypeCollection *hapCol, 
			HapNode *node, unsigned int *haplotype){
  double rv;
  unsigned int i;
  if(haplotype == NULL){
    haplotype = malloc(sizeof(unsigned int) * hapCol->currentLoci);
  }

  rv  = node->totalProb;

  for(i=hapCol->currentLoci; i>0; i--){
    haplotype[i-1] = node->allele;
    node=node->parent;
  }

  return rv;
}

HaplotypeCollection *createHaplotypeCollection(unsigned int nLoci){
  HaplotypeCollection *hapCol;
  unsigned int i;
  hapCol = malloc(sizeof(HaplotypeCollection));
  hapCol->totalLoci = nLoci;
  hapCol->currentLoci = 0;

  hapCol->loci = malloc(sizeof(Locus *)*nLoci);
  for(i=0; i<nLoci; ++i){
    hapCol->loci[i] = NULL;
  }
  hapCol->evidenceCollection.next = NULL;

  hapCol->N = 0;

  return hapCol;
}

void destroyHaplotypeCollection(HaplotypeCollection *hapCol){
  Locus * locus;
  HapNode *node, *next;
  unsigned int i;
  for(i=0; hapCol->loci[i] != NULL && i<hapCol->totalLoci; ++i){
    locus = hapCol->loci[i];
    node = locus->allNodes;
    while(node!=NULL){
      next=node->next;
      freeHapNode(node);
      node=next;
    }
    freeLocus(locus);
  }
  free(hapCol->loci);
  destroyEvidence(hapCol->evidenceCollection.next);
  free(hapCol);
}

void destroyEvidence(ReadEvidence *evidence){
  ReadEvidence *temp;
  while(evidence!=NULL){
    temp = evidence->next;
    free(evidence);
    evidence=temp;
  }
}


/* PartialRead manipulations */
ReadEvidence *createEvidence(HaplotypeCollection *hapCol, PartialRead * read){
  ReadEvidence *rEvidence = malloc(sizeof(ReadEvidence));
  rEvidence->multiplicity = read->multiplicity;

  rEvidence->next = hapCol->evidenceCollection.next;
  hapCol->evidenceCollection.next = rEvidence;
  hapCol->N += read->multiplicity;

  rEvidence->pRead = read;
  return rEvidence;
}

int addRead(HaplotypeCollection *hapCol, PartialRead * read){
  /*
    RETURN VALUES:
    1 -- partial match
    0 -- matched read until the end
    -1 -- hapCol not expanded upto the start locus
    -2 -- read matched no haplotypes
   */
  Locus *locus;
  HapNode *node;
  int allele;
  int retval;
  unsigned int offset;
  PartialRead *pRead = copyRead(read);
  ReadEvidence * newEvidence = createEvidence(hapCol, pRead);
  ReadEvidenceList * nodeEvidence;

  for(offset=0; (offset<pRead->length) && (pRead->alleles[offset]<0); ++offset){
  };

  if(offset+read->start >= hapCol->currentLoci){
    return -1;
  }

  allele= pRead->alleles[offset];

  locus = hapCol->loci[read->start+offset];
  retval = -2;
  pRead->refCount = 0;
  for(node = locus->alleleIndex[allele]; 
      node!=NULL && node->allele == allele; node=node->next){
    nodeEvidence = malloc(sizeof(ReadEvidenceList));
    nodeEvidence->rEvidence=newEvidence;
    nodeEvidence->offset = offset+1;
    pRead->refCount+=matchPath(node,nodeEvidence);  
    retval = 1;
  }
  if(pRead->refCount==0){
    freeRead(newEvidence->pRead);
    newEvidence->pRead=NULL;
    retval = 0;
  }
  return retval;
}

HapNode * matchNextLocus(HapNode *node, unsigned int allele){
  HapNode *child;
  for(child=node->children; child!=NULL; child=child->next_sibling){
    if(child->allele==allele)
      return child;
  }
  return NULL;
}

unsigned int matchPath(HapNode *node, ReadEvidenceList *nodeEvidence){
  HapNode *pathNode;
  ReadEvidenceList *tempEvidence;
  int *alleles = nodeEvidence->rEvidence->pRead->alleles;
  unsigned int length = nodeEvidence->rEvidence->pRead->length;
  unsigned int nhap = 0;
  while(nodeEvidence->offset < length){
    if(alleles[nodeEvidence->offset]>=0){
      pathNode=matchNextLocus(node, alleles[nodeEvidence->offset]);
      if(pathNode==NULL){
	addNodeEvidence(node,nodeEvidence);
	return 1;
      }
      node=pathNode;
      nodeEvidence->offset++;
    }
    else{
      if (node->children==NULL){
	addNodeEvidence(node,nodeEvidence);
	return 1;
      }
      nodeEvidence->offset++;
      for(pathNode = node->children; pathNode!=NULL; pathNode=pathNode->next_sibling){
	tempEvidence = malloc(sizeof(ReadEvidenceList));
	tempEvidence->rEvidence=nodeEvidence->rEvidence;
	tempEvidence->offset = nodeEvidence->offset;
	nhap+=matchPath(pathNode, nodeEvidence);
	nodeEvidence=tempEvidence;
      }
      free(nodeEvidence);
      return nhap;
    }
  }
  addNodeEvidence(node,nodeEvidence);
  return 1;
}

int removeRead(HaplotypeCollection *hapCol, PartialRead *read){
  
  return 0;
}

PartialRead *copyRead(PartialRead *pRead){
  PartialRead * rv = malloc(sizeof(PartialRead));
  if(pRead!=NULL){
    memcpy(rv,pRead,sizeof(PartialRead));
    rv->alleles=malloc(sizeof(unsigned int)*pRead->length);
    memcpy(rv->alleles, pRead->alleles, sizeof(unsigned int)*pRead->length);
  }
  return rv;
}

void freeRead(PartialRead *pRead){
  if(pRead != NULL){
    if(pRead->alleles!=NULL){
      free(pRead->alleles);
    }
    free(pRead);
  }
}

